from exp import data, run
import os, json, random
from datetime import datetime
from com import emb, k_shot
from fuzzywuzzy import process

SKILLS = {
	'Pre-Algebra': ['average calculations', 'ratio and proportion', 'geometry', 'basic arithmetic operations', 'fractions and decimals', 'probability and combinatorics', 'multiplication and division', 'counting and number theory', 'prime number theory', 'multiples and zero properties', 'solving linear equation', 'circles', 'exponentiation rules', 'perimeter and area'],
	'Algebra': ['combinatorial operations and basic arithmetic', 'function skills', 'calculation and conversion skills', 'solving equations', 'inequality skills', 'graph and geometry skills', 'number theory skills', 'factoring skills', 'complex number skills', 'sequence and series skills', 'quadratic equation skills', 'geometric sequence skills', 'polynomial skills', 'ratio and proportion skills', 'logarithmic and exponential skills', 'algebraic manipulation skills', 'distance and midpoint skills', 'arithmetic skills', 'exponent and root skills', 'algebraic expression skills', 'function composition skills'],
	'Inter-Algebra': ['solving inequalities', 'understanding and application of functions', 'inequality solving and understanding', 'calculus optimization skills', 'simplifying and applying floor and ceiling functions', 'function composition and transformation', 'quadratic equations and solutions', 'polynomial skills', 'understand summation and analysis of series', 'sequence and series analysis skills', 'understanding and utilizing infininte series', 'recursive functions and sequences', 'complex number manipulation and operations', 'understanding ellipse properties', 'complex numbers related skills', 'fication and basic operations', 'understanding logarithmic properties and solving equations', 'understanding and manipulation of rational functions', 'properties and application of exponents', 'algebraic manipulation and equations', 'prime number recognition and properties', 'absolute value skills', 'graph understanding and interpretation'],
	'Geometry': ['understanding circle properties and algebraic manipulation', 'other geometric skills', 'pythagorean skills', 'quadrilateral and polygon skills', 'triangle geometry skills', 'calculus skills', '3d geometry and volume calculation skills', 'circle geometry skills', 'area calculation skills', 'coordinate geometry and transformation skills', 'ratio and proportion skills', 'trigonometry skills', 'combinatorics and probability skills', 'algebraic skills'],
	'Number Theory': ['base conversion', 'prime number theory', 'greatest common divisor calculations', 'modular arithmetic', 'solving equations', 'number theory', 'factorization', 'division and remainders', 'exponentiation', 'sequence analysis', 'arithmetic sequences', 'basic arithmetic', 'polynomial operations', 'understanding of fractions', 'number manipulation'],
	'Precalculus': ['matrix operations', 'geometric series comprehension', 'basic trigonometry', 'vector operations', 'coordinate systems', 'trigonometric calculations', 'complex numbers', 'geometric relations', 'calculus', 'algebra and equations', 'three dimensional geometry', 'arithmetic operations', 'parametric equations', 'sequences series and summation', 'geometry triangle properties', 'geometry and space calculation', 'determinant calculation', 'geometry transforms', 'complex number operations'],
	'Probability': ['probability calculation with replacement', 'combinatorics knowledge', 'probability theory and distribution', 'combinatorial mathematics', 'counting principals', 'permutation and combinations', 'probability concepts and calculations', 'calculating and understanding combinations', 'number theory and arithmetic operations', 'factorials and prime factorization', 'understanding and applying combinatorics concepts'],
	}

SKILLS = [v for vs in SKILLS.values() for v in vs]

# question: <question 0> skill: s 0

SELECT_PROMPT_1 = '''
'''

def strip_skill(resp):
	skill = resp.split(', reason:')[0]
	if ':' in skill:
		skill = skill.split(':')[-1]
	skill = skill.strip()
	if skill not in SKILLS:
		skill = process.extract(skill, SKILLS, limit=1)[0][0]
	assert skill in SKILLS, skill
	return skill	

def make_example_reasoner(model_id, question, verbal=True):
	resp = run.generate(model_id, question, 'com/reasoner_prompts', 'skill_select_1',
						skills=SKILLS, 
						question=question.question, 
						solution=question.cot, 
						verbal=verbal,
						)
	return strip_skill(resp), '', []

out_path = '../OUT/skill'

def make_examples():
	math500 = data.get_data('MATH500')
	# task 1
	outfile1 = f'{out_path}/skill_examples_0.jsonl'
	print(f'# {outfile1} {datetime.now()}')
	run.run_exp(math500, make_example_reasoner, num_samples=8, max_retries=50, outfile=outfile1)
	# task 2
	outfile2 = f'{out_path}/skill_examples.json'
	examples_0 = emb.load_jsonl(outfile1)
	skills = {s:{} for s in SKILLS}
	for e in examples_0:
		pid = e['pid']
		skill = e['pred'][1]
		if pid not in skills[skill]:
			skills[skill][pid] = 0
		skills[skill][pid] += 1
	examples = {}
	problems = {e['pid']:e for e in math500}
	for s,pc in skills.items():
		pc = [[k,v] for k,v in pc.items()]
		pc.sort(key=lambda kv: kv[1], reverse=True)
		examples[s] = [problems[pid] for pid,c in pc]
	with open(outfile2, 'w') as fp:
		json.dump(examples, fp, indent='\t', ensure_ascii=False)

######################

class Skill_Reasoner:
	def __init__(self, examples, k=8):
		self.examples = examples
		self.skill_examples = []
		for s,es in examples.items():
			if len(es) > 0:
				problem = es[0]['problem'].strip().replace('\n', '  ')
				self.skill_examples.append(f'question: {json.dumps(problem)} skill: {json.dumps(s)}')
		self.skill_examples = '\n'.join(self.skill_examples)
		self.k = k

	def reason(self, model_id, question, verbal=True):
		resp = run.generate(model_id, question, 'com/reasoner_prompts', 'skill_select_2',
							skills=SKILLS, 
							question=question.question, 
							examples=self.skill_examples, 
							verbal=verbal,
							)
		skill = strip_skill(resp)
		examples = self.examples[skill]
		random.shuffle(examples)
		examples = [e for e in examples[:self.k+1] if e['problem'] != question]
		examples = k_shot._examples_repr(examples[:self.k])
		resp = run.generate(model_id, question, 'com/reasoner_prompts', 'k_shot',
							question=question.question, 
							examples=examples, 
							verbal=verbal,
							)
		return question.extra_asnwer(resp), 'The answer is', []

######################

def test(test_set):
	skill_examples_path = f'{out_path}/skill_examples.json'
	with open(skill_examples_path) as fp:
		skill_examples = json.load(fp)
	test_examples = data.get_data(test_set)
	reasoner = Skill_Reasoner(skill_examples)
	for e in test_examples[:10]:
		run.test(e, reasoner.reason)

def exp(test_set):
	skill_examples_path = f'{out_path}/skill_examples.json'
	with open(skill_examples_path) as fp:
		skill_examples = json.load(fp)
	test_examples = data.get_data(test_set)
	reasoner = Skill_Reasoner(skill_examples)
	outfile = f'{out_path}/{test_set}_skill.jsonl'
	print(f'# {outfile} {datetime.now()}')
	run.run_exp(test_examples, reasoner.reason,  num_samples=8, max_retries=8, outfile=outfile)

def main():
	if not os.path.isdir(out_path):
		os.mkdir(out_path)
	make_examples()
	# test('MATH500')
	exp('MATH500')

if __name__ == '__main__':
	main()
